# -*- coding:utf-8 -*-
import numpy as np
import os
from tensorflow import keras
from PIL import Image
from sklearn.model_selection import train_test_split
import cv2
dimen = 32
from opencv import cv_show
from conneted_image import get_largest_conneted_area, get_largest_contour, opening

def DataConverter():
    dir_path = 'DR_vertebral/'
    output_path = 'processed_data/'
    sub_dir_list = os.listdir(dir_path)
    images = list()
    labels = list()
    for i in range(len(sub_dir_list)):
        label = i
        image_names = os.listdir(dir_path + sub_dir_list[i])
        for image_path in image_names:
            path = dir_path + sub_dir_list[i] + "/" + image_path
            image = Image.open(path).convert('L')
            resize_image = image.resize((dimen, dimen))
            array = list()
            for x in range(dimen):
                sub_array = list()
                for y in range(dimen):
                    sub_array.append(resize_image.load()[x, y])
                array.append(sub_array)
            image_data = np.array(array)
            image = np.array(np.reshape(image_data, (dimen, dimen, 1))) / 255
            images.append(image)
            labels.append(label)

    x = np.array(images)
    y = np.array(keras.utils.to_categorical(np.array(labels), num_classes=len(sub_dir_list)))

    train_features, test_features, train_labels, test_labels = train_test_split(x, y, test_size=0.2)

    np.save('{}x.npy'.format(output_path), train_features)
    np.save('{}y.npy'.format(output_path), train_labels)
    np.save('{}test_x.npy'.format(output_path), test_features)
    np.save('{}test_y.npy'.format(output_path), test_labels)


def segment_largest_roi_from_image(image:np.ndarray, debug=False):
    mask:np.ndarray = np.uint8(image > 0) * 255

    # retval, mask = cv2.threshold(image, 0, 255, cv2.THRESH_OTSU)

    h, w = mask.shape
    # if debug:
    #     cv_show('mask', mask, resize=0.7)
    #     cv2.waitKey(0)

    # erode_mask = erode(mask, ks=100, iters=10)
    dilate_mask = opening(mask, ks=20, iters=2)
    if debug:
        cv_show('dilate_mask', np.concatenate((image, mask, dilate_mask), axis=1), resize=0.7)
        cv2.waitKey(0)

    contours, hierarchy = cv2.findContours(dilate_mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)  # 轮廓提取
    dilate_min_x, dilate_min_y, dilate_max_x, dilate_max_y = get_largest_conneted_area(contours)
    roi_xyxy =  np.array([dilate_min_x, dilate_min_y, dilate_max_x, dilate_max_y])

    roi_xyxy[0::2] = np.clip(roi_xyxy[0::2], 0, w)
    roi_xyxy[1::2] = np.clip(roi_xyxy[1::2], 0, h)

    roi_xywh = np.array([roi_xyxy[0], roi_xyxy[1], roi_xyxy[2]-roi_xyxy[0], roi_xyxy[3]-roi_xyxy[1]])

    roi_w_ratio = roi_xywh[2] / w
    roi_h_ratio = roi_xywh[3] / h

    defalut_offset_w = 10
    defalut_offset_h = 10
    if roi_h_ratio < 0.8:
        defalut_offset_h = int((h - roi_xywh[3]) / 4) +50

    roi_xyxy[0] = roi_xyxy[0] - defalut_offset_w
    roi_xyxy[1] = roi_xyxy[1] - defalut_offset_h
    roi_xyxy[2] = roi_xyxy[2] + defalut_offset_w
    roi_xyxy[3] = roi_xyxy[3] + defalut_offset_h

    roi_xyxy[0::2] = np.clip(roi_xyxy[0::2], 0, w)
    roi_xyxy[1::2] = np.clip(roi_xyxy[1::2], 0, h)

    return roi_xyxy.tolist()


def crop_image():
    debug = False
    dir_path = 'DR_vertebral/'
    output_path = 'DR/crop_images/'
    sub_dir_list = os.listdir(dir_path)
    for i in range(len(sub_dir_list)):
        image_names = os.listdir(dir_path + sub_dir_list[i])
        for image_path in image_names:
            path = dir_path + sub_dir_list[i] + "/" + image_path
            image = cv2.imread(path)
            # crop images
            COLOR_BGR2GRAY = 6
            gray_image = cv2.cvtColor(image, code=COLOR_BGR2GRAY)
            # 1.crop roi
            min_x, min_y, max_x, max_y = segment_largest_roi_from_image(gray_image, debug=debug)
            crop_image = image[min_y:max_y, min_x:max_x]
            if debug:
                cv2.imshow('crop', crop_image)
                cv2.waitKey(0)
            # save iamge
            cv2.imwrite(output_path + image_path, crop_image)


if __name__ == "__main__":
    # crop_image()
    DataConverter()